{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Natural Language Question Answering\n",
    "\n",
    "### Anna Bethke, Andy Keller\n",
    "### Originally Presented at Intel AI DevCon 2018\n",
    "\n",
    "This notebook gives an overview of an End-to-End Memory Network for goal oriented dialogue, implemented in ngraph.\n",
    "\n",
    "Goal oriented dialogue is a subset of open-domain dialogue where an automated agent has a specific\n",
    "goal for the outcome of the interaction. At a high level, the system needs to understand a user\n",
    "request and complete a related task with a clear goal within a limited number of dialog turns.\n",
    "This task could be making a restaurant reservation, placing an order, setting a timer, or many of the digital personal assistant tasks.\n",
    "\n",
    "End-to-End Memory Networks are generic semi-recurrent neural networks which allow for a bank of\n",
    "external memories to be read from and used during execution. They can be used in place of traditional\n",
    "slot-filling algorithms to accomplish goal oriented dialogue tasks without the need for expensive\n",
    "hand-labeled dialogue data. End-to-End Memory Networks have also been shown to be useful for\n",
    "Question-Answering and information retrieval tasks.\n",
    "\n",
    "This demonstration will go through a short training cycle to conduct full dialogues, and an interactive demonstration of what the network can do with a fully trained model. Training a full model generally does not take a significant amount of time, but it may be longer than the duration of this session. The implementation is based off the paper by A. Bordes, Y. Boureau, J. Weston. [Learning End-to-End Goal-Oriented Dialog](https://arxiv.org/abs/1605.07683) 2016 and the Github repository [chatbot-MemN2N-tensorflow](https://github.com/vyraun/chatbot-MemN2N-tensorflow)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The model was trained and evaluated on the 6 bAbI Dialog tasks with the following results.\n",
    "\n",
    "| Task | This  | Published |  This (w/ match-type) | Published (w/ match-type)|\n",
    "|------|--------|-----------| ---------------------|--------------------------|\n",
    "| 1    | 99.8   | 99.9      | 100.0                | 100.0                    |\n",
    "| 2    | 100.0  | 100.0     | 100.0                | 98.3                     |\n",
    "| 3    | 74.8   | 74.9      | 74.6                 | 74.9                     |\n",
    "| 4    | 57.2   | 59.5      | 100.0                | 100.0                    |\n",
    "| 5    | 96.4   | 96.1      | 95.6                 | 93.4                     |\n",
    "| 6    | 48.1   | 41.1      | 45.4                 | 41.0                     |"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![](https://camo.githubusercontent.com/ba1c7dbbccc5dd51d4a76cc6ef849bca65a9bf4d/687474703a2f2f692e696d6775722e636f6d2f6e7638394a4c632e706e67)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "![](https://i.imgur.com/5pQJqjM.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our first step is to import our necessary libraries. For simplicities sake, in this tutorial, we will be using some code from the NLP Architect github reposiotry. The code presented in this demonstration notebook is also included in that repository. \n",
    "\n",
    "NLP Architect is a repository for models exploring the state of the art deep learning techniques for natural language processing and natural language understanding. It is intended to be a platform for future research and collaboration.\n",
    "\n",
    "The library includes our past and ongoing NLP research efforts as part of Intel AI Lab. For more information on the library please see the documentation at https://intellabs.github.io/nlp-architect, or the github repository at  https://github.com/IntelLabs/nlp-architect."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from __future__ import division\n",
    "from __future__ import print_function\n",
    "from __future__ import unicode_literals\n",
    "from __future__ import absolute_import\n",
    "from builtins import input\n",
    "import numpy as np\n",
    "from copy import copy\n",
    "from functools import reduce\n",
    "import itertools\n",
    "from tqdm import tqdm\n",
    "\n",
    "import ngraph as ng\n",
    "from ngraph.frontends.neon import Layer\n",
    "from ngraph.frontends.neon import GaussianInit\n",
    "from ngraph.frontends.neon import make_bound_computation\n",
    "from ngraph.frontends.neon import NgraphArgparser\n",
    "from ngraph.frontends.neon import ArrayIterator\n",
    "from ngraph.frontends.neon import Saver\n",
    "from ngraph.frontends.neon import GaussianInit, Adam, GradientDescentMomentum, RMSProp\n",
    "import ngraph.transformers as ngt\n",
    "\n",
    "from nlp_architect.data.babi_dialog import BABI_Dialog\n",
    "from nlp_architect.models.memn2n_dialogue import MemN2N_Dialog\n",
    "from nlp_architect.contrib.ngraph.modified_lookup_table import ModifiedLookupTable\n",
    "from contextlib import closing\n",
    "\n",
    "from ipywidgets import interact, interactive, fixed, interact_manual\n",
    "import ipywidgets as widgets\n",
    "\n",
    "from IPython.display import display\n",
    "import time"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting up arguments and getting the data\n",
    "\n",
    "To start with, we will set some of the parameters we will need throughout this demonstration. Feel free to change these around. The most important thing to remember though is that when you want to run the interactive demo, you need to have the same parameters that you used when you saved a particular model. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "task = 5 # The task ID to train/test on from bAbI-dialog dataset (1-6)\n",
    "emb_size = 32 # Size of the word-embedding used in the model.\n",
    "nhops = 3 # Number of memory hops in the network\n",
    "use_match_type = False # Use match type features\n",
    "cache_match_type = False # Cache pre-processed match type answers\n",
    "use_oov = False # Use OOV test set and for the interactive mode, allow out of vocabulary words\n",
    "lr = 0.001 # learning rate\n",
    "eps = 1e-8 # Epsilon used to avoid divide by zero in softmax renormalization\n",
    "model_file = './memn2n_task5_weights_retrain.npz' # Where you want to store your model weights\n",
    "data_dir = './data/' # The directory of the dataset\n",
    "restore = False # Restore model weights if found\n",
    "batch_size = 32 # Size of the batch during training\n",
    "gradient_clip_norm = 40.0 # Clip gradients such that norm is below this value\n",
    "cache_vectorized = False # Cache preprocessed/vectorized dataset"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Before we dive into the memory network code, lets get to know our data a little bit better. The dataset used for training and evaluation is under the umbrella of the Facebook bAbI dialog tasks\n",
    "(https://research.fb.com/downloads/babi/). The terms and conditions of the data set license apply. Intel does not grant any rights to the data files. The dataset is automatically downloaded if not found,\n",
    "and the preprocessing all happens at the beginning of training.\n",
    "\n",
    "There are six separate tasks, tasks 1 through 5 are from simulated conversations between a customer\n",
    "and a restaurant booking bot (created by Facebook), and task 6 is more realistic natural language\n",
    "restaurant booking conversations as part of the `dialog state tracking challenge`.\n",
    "\n",
    "The descriptions of the six tasks are as follow:\n",
    "\n",
    "- bAbI dialog dataset:\n",
    "    - Task 1: Issuing API Calls\n",
    "    - Task 2: Updating API Calls\n",
    "    - Task 3: Displaying Options\n",
    "    - Task 4: Providing Extra Information\n",
    "    - Task 5: Conducting Full Dialogs\n",
    "\n",
    "- Dialog State Tracking Challenge 2 Dataset:\n",
    "    - Task 6: DSTC2 Full Dialogs\n",
    "    \n",
    "We have written a function BABI_Dialog which is able to take a directory of the dataset and with a few parameters prepare the data for the memory network task. Lets look to see what our data looks like before and after it goes through the this class. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The core of the dataset is conversations, for example the first conversation for Task 5 (our focus today) is:\n",
    "\n",
    "```\n",
    "1 good morning\thello what can i help you with today\n",
    "2 i'd like to book a table with italian food\ti'm on it\n",
    "3 <SILENCE>\twhere should it be\n",
    "4 in paris\thow many people would be in your party\n",
    "5 for six people please\twhich price range are looking for\n",
    "6 in a cheap price range please\tok let me look into some options for you\n",
    "7 <SILENCE>\tapi_call italian paris six cheap\n",
    "8 actually i would prefer for two people\tsure is there anything else to update\n",
    "9 instead could it be in madrid\tsure is there anything else to update\n",
    "10 instead could it be with spanish food\tsure is there anything else to update\n",
    "11 no\tok let me look into some options for you\n",
    "12 <SILENCE>\tapi_call spanish madrid two cheap\n",
    "13 resto_madrid_cheap_spanish_1stars R_phone resto_madrid_cheap_spanish_1stars_phone\n",
    "14 resto_madrid_cheap_spanish_1stars R_cuisine spanish\n",
    "15 resto_madrid_cheap_spanish_1stars R_address resto_madrid_cheap_spanish_1stars_address\n",
    "16 resto_madrid_cheap_spanish_1stars R_location madrid\n",
    "17 resto_madrid_cheap_spanish_1stars R_number two\n",
    "18 resto_madrid_cheap_spanish_1stars R_price cheap\n",
    "19 resto_madrid_cheap_spanish_1stars R_rating 1\n",
    "20 resto_madrid_cheap_spanish_6stars R_phone resto_madrid_cheap_spanish_6stars_phone\n",
    "21 resto_madrid_cheap_spanish_6stars R_cuisine spanish\n",
    "22 resto_madrid_cheap_spanish_6stars R_address resto_madrid_cheap_spanish_6stars_address\n",
    "23 resto_madrid_cheap_spanish_6stars R_location madrid\n",
    "24 resto_madrid_cheap_spanish_6stars R_number two\n",
    "25 resto_madrid_cheap_spanish_6stars R_price cheap\n",
    "26 resto_madrid_cheap_spanish_6stars R_rating 6\n",
    "27 resto_madrid_cheap_spanish_8stars R_phone resto_madrid_cheap_spanish_8stars_phone\n",
    "28 resto_madrid_cheap_spanish_8stars R_cuisine spanish\n",
    "29 resto_madrid_cheap_spanish_8stars R_address resto_madrid_cheap_spanish_8stars_address\n",
    "30 resto_madrid_cheap_spanish_8stars R_location madrid\n",
    "31 resto_madrid_cheap_spanish_8stars R_number two\n",
    "32 resto_madrid_cheap_spanish_8stars R_price cheap\n",
    "33 resto_madrid_cheap_spanish_8stars R_rating 8\n",
    "34 <SILENCE>\twhat do you think of this option: resto_madrid_cheap_spanish_8stars\n",
    "35 no this does not work for me\tsure let me find an other option for you\n",
    "36 <SILENCE>\twhat do you think of this option: resto_madrid_cheap_spanish_6stars\n",
    "37 do you have something else\tsure let me find an other option for you\n",
    "38 <SILENCE>\twhat do you think of this option: resto_madrid_cheap_spanish_1stars\n",
    "39 it's perfect\tgreat let me do the reservation\n",
    "40 may i have the phone number of the restaurant\there it is resto_madrid_cheap_spanish_1stars_phone\n",
    "41 thanks\tis there anything i can help you with\n",
    "42 no thank you\tyou're welcome\n",
    "```\n",
    "\n",
    "It is a mixture of conversation between the user and agent along with API calls, and the resulting database responses, from the agent.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now lets process the data\n",
    "# This involves parsing the above into user/bot utterances\n",
    "# Then encoding each word using an integer index into our vocabulary.\n",
    "# This does take some time (5-10 minutes)\n",
    "babi = BABI_Dialog(\n",
    "    path=data_dir,\n",
    "    task=task,\n",
    "    oov=use_oov,\n",
    "    use_match_type=use_match_type,\n",
    "    cache_match_type=cache_match_type)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The data within the babi object is a dictionary broken up by train/test\n",
    "\n",
    "If we look at the training data there are answers, utterances, a memory mask and the memory. Each of the data objects has axes information along with the actual data.\n",
    "The data has been preprocessed to be a bag of words representation. The data is already fairly clean in that everything is already lowercase and there are no spelling or gramatical errors."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "babi.data_dict['train'].keys()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "babi.data_dict['train']['user_utt']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# So for example the second sentence \"i'd like to book a table with italian food\"  looks like:\n",
    "babi.data_dict['train']['user_utt']['data'][1]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The memory contains information from of the conversation thus far\n",
    "# so after one back and forth with the agent, it looks like the following:\n",
    "babi.data_dict['train']['memory']['data'][1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We can see that all utterances are padded to the maximum sentence length with 0's, and the empty memory slots are also filled with all 0 sentences. These values are eplicitly forced to zero in our embedding layer, and never updated."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "babi.data_dict['train']['memory_mask']['data'][1]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The 'memory_mask' is used to mask out empty memory slots after the softmax is applied during the computation of the scalar similarity values. This is a trick to ensure model correctness, and emulate a softmax over a dynamic memory size. "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setting up for training\n",
    "\n",
    "Now that we have the data, lets look at the model we will use to train it\n",
    "\n",
    "The memory network heavily uses lookup tables - these are mappings from one object space to another. They are being used to map our vocab-size one-hot vectors to embedding-size dense vectors, allowing the model to decrease the input dimensionality and additionally learn how words are associated with each other. \n",
    "\n",
    "The class MemN2N_Dialog is the core of the memory network model. It starts with the various arguments that we will pass in when training is initialized."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class MemN2N_Dialog(Layer):\n",
    "    \"\"\"\n",
    "    End-to-End Memory Networks for Goal Oriented Dialogue\n",
    "\n",
    "    After the model is initialized, it accepts a BABI_Dialog class formatted dataset \n",
    "    as input and returns a probability distribution over candidate answers. \n",
    "\n",
    "    Args:\n",
    "        cands (np.array): Vectorized array of potential candidate answers, encoded\n",
    "            as integers, as returned by BABI_Dialog class. Shape = [num_cands, max_cand_length]\n",
    "        num_cands (int): Number of potential candidate answers. \n",
    "        max_cand_len (int): Maximum length of a candidate answer sentence in number of words. \n",
    "        memory_size (int): Maximum number of sentences to keep in memory at any given time.\n",
    "        max_utt_len (int): Maximum length of any given sentence / user utterance \n",
    "        vocab_size (int): Number of unique words in the vocabulary + 2 (0 is reserved for \n",
    "            a padding symbol, and 1 is reserved for OOV)\n",
    "        emb_size (int): Dimensionality of word embeddings to use \n",
    "        batch_size (int): Number of training examples per batch \n",
    "        use_match_type (bool, optional): Flag to use match-type features\n",
    "        kb_ents_to_type (dict, optional): For use with match-type features, dictionary of \n",
    "            entities found in the dataset mapping to their associated match-type\n",
    "        kb_ents_to_cand_idxs (dict, optional): For use with match-type features, dictionary\n",
    "            mapping from each entity in the  knowledge base to the set of indicies in the\n",
    "            candidate_answers array that contain that entity.\n",
    "        match_type_idxs (dict, optional): For use with match-type features, dictionary \n",
    "            mapping from match-type to the associated fixed index of the candidate vector\n",
    "            which indicated this match type.\n",
    "        nhop (int, optional): Number of memory-hops to perform during fprop \n",
    "        eps (float, optional): Small epsilon used for numerical stability in \n",
    "            softmax renormalization\n",
    "        init (Initalizer, optional): Initalizer object used to initialize lookup table\n",
    "            and projection layer.\n",
    "    \"\"\"\n",
    "    def __init__(\n",
    "        self,\n",
    "        cands,\n",
    "        num_cands,\n",
    "        max_cand_len,\n",
    "        memory_size,\n",
    "        max_utt_len,\n",
    "        vocab_size,\n",
    "        emb_size,\n",
    "        batch_size,\n",
    "        use_match_type=False,\n",
    "        kb_ents_to_type=None,\n",
    "        kb_ents_to_cand_idxs=None,\n",
    "        match_type_idxs=None,\n",
    "        nhops=3,\n",
    "        eps=1e-6,\n",
    "        init=GaussianInit(\n",
    "            mean=0.0,\n",
    "            std=0.1)):\n",
    "        super(MemN2N_Dialog, self).__init__()\n",
    "\n",
    "        self.cands = cands\n",
    "        self.memory_size = memory_size\n",
    "        self.max_utt_len = max_utt_len\n",
    "        self.vocab_size = vocab_size\n",
    "        self.num_cands = num_cands\n",
    "        self.max_cand_len = max_cand_len\n",
    "        self.batch_size = batch_size\n",
    "        self.use_match_type = use_match_type\n",
    "        self.kb_ents_to_type = kb_ents_to_type\n",
    "        self.kb_ents_to_cand_idxs = kb_ents_to_cand_idxs\n",
    "        self.match_type_idxs = match_type_idxs\n",
    "        self.nhops = nhops\n",
    "        self.eps = eps\n",
    "        self.init = init\n",
    "\n",
    "        # Make axes\n",
    "        self.batch_axis = ng.make_axis(length=batch_size, name='N')\n",
    "        self.sentence_rec_axis = ng.make_axis(length=max_utt_len, name='REC')\n",
    "        self.memory_axis = ng.make_axis(length=memory_size, name='memory_axis')\n",
    "        self.embedding_axis = ng.make_axis(length=emb_size, name='F')\n",
    "        self.embedding_axis_proj = ng.make_axis(length=emb_size, name='F_proj')\n",
    "        self.cand_axis = ng.make_axis(length=num_cands, name='cand_axis')\n",
    "        self.cand_rec_axis = ng.make_axis(length=max_cand_len, name='REC')\n",
    "\n",
    "        # Weight sharing of A's accross all hops input and output\n",
    "        self.LUT_A = ModifiedLookupTable(\n",
    "            vocab_size, emb_size, init, update=True, pad_idx=0)\n",
    "        # Use lookuptable W to embed the candidate answers\n",
    "        self.LUT_W = ModifiedLookupTable(\n",
    "            vocab_size, emb_size, init, update=True, pad_idx=0)\n",
    "\n",
    "        # Initialize projection matrix between internal model states\n",
    "        self.R_proj = ng.variable(axes=[self.embedding_axis, self.embedding_axis_proj],\n",
    "            initial_value=init)\n",
    "\n",
    "        if not self.use_match_type:\n",
    "            # Initialize constant matrix of all candidate answers\n",
    "            self.cands_mat = ng.constant(self.cands, axes=[self.cand_axis, self.cand_rec_axis])\n",
    "\n",
    "    def __call__(self, inputs):\n",
    "        query = ng.cast_axes(inputs['user_utt'], [self.batch_axis, self.sentence_rec_axis])\n",
    "\n",
    "        # Query embedding [batch, sentence_axis, F]\n",
    "        q_emb = self.LUT_A(query)\n",
    "\n",
    "        # Sum the embeddings\n",
    "        u_0 = ng.sum(q_emb, reduction_axes=[self.sentence_rec_axis])\n",
    "\n",
    "        # Start a list of the internal states of the model. Will be appended to\n",
    "        # after each memory hop\n",
    "        u = [u_0]\n",
    "\n",
    "        for hopn in range(self.nhops):\n",
    "            story = ng.cast_axes(inputs['memory'], [self.batch_axis, self.memory_axis, self.sentence_rec_axis])\n",
    "\n",
    "            # Re-use the query embedding matrix to embed the memory sentences\n",
    "            # [batch, memory_axis, sentence_axis, F]\n",
    "            m_emb_A = self.LUT_A(story)\n",
    "            m_A = ng.sum(m_emb_A, reduction_axes=[self.sentence_rec_axis])  # [batch, memory_axis, F]\n",
    "\n",
    "            # Compute scalar similarity between internal state and each memory\n",
    "            # Equivalent to dot product between u[-1] and each memory in m_A\n",
    "            # [batch, memory_axis]\n",
    "            dotted = ng.sum(u[-1] * m_A, reduction_axes=[self.embedding_axis])\n",
    "\n",
    "            # [batch, memory_axis]\n",
    "            probs = ng.softmax(dotted, self.memory_axis)\n",
    "\n",
    "            # Renormalize probabilites according to non-empty memories\n",
    "            probs_masked = probs * inputs['memory_mask']\n",
    "            renorm_sum = ng.sum(probs_masked, reduction_axes=[self.memory_axis]) + self.eps\n",
    "            probs_renorm = (probs_masked + self.eps) / renorm_sum\n",
    "\n",
    "            # Compute weighted sum of memory embeddings\n",
    "            o_k = ng.sum(probs_renorm * m_A, reduction_axes=[self.memory_axis])  # [batch, F]\n",
    "\n",
    "            # Add the output back into the internal state and project\n",
    "            u_k = ng.cast_axes(ng.dot(self.R_proj, o_k), [\n",
    "                               self.embedding_axis, self.batch_axis]) + u[-1]  # [batch, F_proj]\n",
    "\n",
    "            # Add new internal state\n",
    "            u.append(u_k)\n",
    "\n",
    "        if self.use_match_type:\n",
    "            # [batch_axis, cand_axis, cand_rec_axis, F]\n",
    "            self.cands_mat = inputs['cands_mat']\n",
    "\n",
    "        # Embed all candidate responses using LUT_W\n",
    "        # [<batch_axis>, cand_axis, cand_rec_axis, F]\n",
    "        cand_emb_W = self.LUT_W(self.cands_mat)\n",
    "        # No position encoding added yet\n",
    "        cands_mat_emb = ng.sum(cand_emb_W, reduction_axes=[self.cand_rec_axis])  # [<batch_axis>, cand_axis, F]\n",
    "\n",
    "        # Compute predicted answer from product of final internal state\n",
    "        # and embedded candidate answers\n",
    "        # same as: a_logits = ng.dot(cands_mat_emb, u[-1]) # [batch, cand_axis]\n",
    "        a_logits = ng.sum(u[-1] * cands_mat_emb, reduction_axes=[self.embedding_axis])\n",
    "\n",
    "        # rename V to cand_axis to match answer\n",
    "        a_logits = ng.cast_axes(a_logits, [self.batch_axis, self.cand_axis])\n",
    "        a_pred = ng.softmax(a_logits, self.cand_axis)\n",
    "\n",
    "        return a_pred, probs_renorm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The next stage in our training process is to initialize the memory network model by calling MemN2N_Dialog(). Here is a good place that you can modify some of the hyperparameters as well as other parameters like the initial default initialization. You can also choose a different optimizier. We are using Adam here, but the neon frontend of NGraph also includes GradientDescentMomentum, and RMSprop. These are already imported in case you would like to experiment with them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "weights_save_path = model_file\n",
    "\n",
    "weight_saver = Saver()\n",
    "\n",
    "# Set num iterations to 1 epoch since we loop over epochs & shuffle\n",
    "ndata = babi.data_dict['train']['memory']['data'].shape[0]\n",
    "num_iterations = ndata // batch_size\n",
    "\n",
    "train_set = ArrayIterator(babi.data_dict['train'], batch_size=batch_size,\n",
    "                          total_iterations=num_iterations)\n",
    "dev_set = ArrayIterator(babi.data_dict['dev'], batch_size=batch_size)\n",
    "test_set = ArrayIterator(babi.data_dict['test'], batch_size=batch_size)\n",
    "inputs = train_set.make_placeholders()\n",
    "\n",
    "memn2n = MemN2N_Dialog(\n",
    "    babi.cands,\n",
    "    babi.num_cands,\n",
    "    babi.max_cand_len,\n",
    "    babi.memory_size,\n",
    "    babi.max_utt_len,\n",
    "    babi.vocab_size,\n",
    "    emb_size,\n",
    "    batch_size,\n",
    "    use_match_type=use_match_type,\n",
    "    kb_ents_to_type=babi.kb_ents_to_type,\n",
    "    kb_ents_to_cand_idxs=babi.kb_ents_to_cand_idxs,\n",
    "    match_type_idxs=babi.match_type_idxs,\n",
    "    nhops=nhops,\n",
    "    eps=eps,\n",
    "    init=GaussianInit(\n",
    "        mean=0.0,\n",
    "        std=0.1))\n",
    "\n",
    "# Compute answer predictions\n",
    "a_pred, attention = memn2n(inputs)\n",
    "\n",
    "# specify loss function, calculate loss and update weights\n",
    "loss = ng.cross_entropy_multi(a_pred, inputs['answer'], usebits=True)\n",
    "\n",
    "# Compute sum of losses\n",
    "mean_cost = ng.sum(loss, out_axes=[])\n",
    "\n",
    "# initialize optimizer and create update-op\n",
    "optimizer = Adam(learning_rate=lr)\n",
    "updates = optimizer(loss)\n",
    "batch_cost = ng.sequential([updates, mean_cost])\n",
    "\n",
    "# provide outputs for bound computation\n",
    "train_outputs = dict(batch_cost=batch_cost, train_preds=a_pred)\n",
    "\n",
    "# Create additional inference graphs to handle things like dropout automatically\n",
    "with Layer.inference_mode_on():\n",
    "    a_pred_inference, attention_inference = memn2n(inputs)\n",
    "    eval_loss = ng.cross_entropy_multi(\n",
    "        a_pred_inference, inputs['answer'], usebits=True)\n",
    "\n",
    "# Create interactive and eval outputs\n",
    "interactive_outputs = dict(\n",
    "    test_preds=a_pred_inference,\n",
    "    attention=attention_inference)\n",
    "eval_outputs = dict(test_cross_ent_loss=eval_loss, test_preds=a_pred_inference)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Our last and final stage is to initialize the training loop and run training. This cell will create the computational objects, restore any weights that were previously saved, and then for each epoch, run the memory network model. At the end of each training epoch, the validation set will be evaluated to get a continuous metric on how well the network is being trained."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "epochs = 3\n",
    "save_epochs = 1\n",
    "\n",
    "# Train Loop\n",
    "with closing(ngt.make_transformer()) as transformer:\n",
    "    # bind the computations\n",
    "    train_computation = make_bound_computation(transformer, train_outputs, inputs)\n",
    "    loss_computation = make_bound_computation(transformer, eval_outputs, inputs)\n",
    "    interactive_computation = make_bound_computation(transformer, interactive_outputs, inputs)\n",
    "\n",
    "    weight_saver.setup_save(transformer=transformer, computation=train_outputs)\n",
    "\n",
    "    # Load weights if found\n",
    "    if restore and os.path.exists(weights_save_path):\n",
    "        print(\"Loading weights from {}\".format(weights_save_path))\n",
    "        weight_saver.setup_restore(\n",
    "            transformer=transformer,\n",
    "            computation=train_outputs,\n",
    "            filename=weights_save_path)\n",
    "        weight_saver.restore()\n",
    "    elif restore and os.path.exists(weights_save_path) is False:\n",
    "        print(\"Could not find weights at {}. \".format(weights_save_path)\n",
    "              + \"Running with random initialization.\")\n",
    "\n",
    "    for e in range(epochs):\n",
    "        train_error = []\n",
    "        train_cost = []\n",
    "        \n",
    "        # Loop over training batches, and update model with the train computation\n",
    "        for idx, data in enumerate(\n",
    "            tqdm(train_set, total=train_set.nbatches,\n",
    "                 unit='minibatches', desc=\"Epoch {}\".format(e))):\n",
    "            train_output = train_computation(data)\n",
    "            \n",
    "            # Store loss and error for monitoring\n",
    "            train_cost.append(train_output['batch_cost'])\n",
    "            preds = np.argmax(train_output['train_preds'], axis=1)\n",
    "            error = np.mean(data['answer'].argmax(axis=1) != preds)\n",
    "            train_error.append(error)\n",
    "\n",
    "        train_cost_str = \"Epoch {}: train_cost {}, train_error {}\".format(\n",
    "            e, np.mean(train_cost), np.mean(train_error))\n",
    "        print(train_cost_str)\n",
    "\n",
    "        if e % save_epochs == 0:\n",
    "            print(\"Saving model to {}\".format(weights_save_path))\n",
    "            weight_saver.save(filename=weights_save_path)\n",
    "            print(\"Saving complete - Running validation set\")\n",
    "\n",
    "        # Eval after each epoch\n",
    "        test_loss = []\n",
    "        test_error = []\n",
    "        for idx, data in enumerate(dev_set):\n",
    "            # only compute loss on evaluation set\n",
    "            test_output = loss_computation(data)\n",
    "            test_loss.append(np.sum(test_output['test_cross_ent_loss']))\n",
    "            preds = np.argmax(test_output['test_preds'], axis=1)\n",
    "            error = np.mean(data['answer'].argmax(axis=1) != preds)\n",
    "            test_error.append(error)\n",
    "\n",
    "        val_cost_str = \"Epoch {}: validation_cost {}, validation_error {}\".format(\n",
    "            e, np.mean(test_loss), np.mean(test_error))\n",
    "        print(val_cost_str)\n",
    "\n",
    "        # Shuffle training set and reset the others\n",
    "        shuf_idx = np.random.permutation(range(train_set.data_arrays['memory'].shape[0]))\n",
    "        train_set.data_arrays = {k: v[shuf_idx] for k, v in train_set.data_arrays.items()}\n",
    "        train_set.reset()\n",
    "        dev_set.reset()\n",
    "\n",
    "    print('Training Complete.')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Moving to interactive.py"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now that we have a trained model we can talk to it! For that though it is best to go to the command line!\n",
    "\n",
    "In the command line type:\n",
    "\n",
    "```cd private-nlp-architect\n",
    "python interactive.py --task 5 --data_dir ./data/ --model_file ./examples/memn2n_dialogue/memn2n_weights.npz```\n",
    "\n",
    "Now it's up to you to play with the model architecture above and hyperparameter settings to observe changes in model performance!"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "nlp_architect_env",
   "language": "python",
   "name": "nlp_architect_env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
